#include "CoordinateProcess.h"
#include "../Component/CoordinateModel.h"
#include "../../Render/Buffer/VertexBuffer.h"
#include "../../Render/Buffer/IndexBuffer.h"
#include "../../Render/Buffer/CustomBuffer.h"

CoordinateProcess::CoordinateProcess()
{
	mColorProcess = PREV_PROCESS;
	mColorName = "ColorRenderTarget";
	mColorDirty = false;
}

CoordinateProcess::CoordinateProcess(const std::string& colorProcess)
{
	mColorProcess = std::move(colorProcess);
	mColorName = "ColorRenderTarget";
	mColorDirty = false;
}

CoordinateProcess::CoordinateProcess(const std::string& colorProcess, const std::string& colorName)
{
	mColorProcess = std::move(colorProcess);
	mColorName = colorName;
	mColorDirty = false;
}

void CoordinateProcess::OnInit()
{
	mCoordinateEntity = mGameScene->CreateEntity<Entity>();
	CoordinateModel* coordinateModel = mCoordinateEntity->AddComponent<CoordinateModel>();
	float screenX = 0.1f * GRenderState->GetHeight() / GRenderState->GetWidth();
	mCoordinateEntity->GetTransform().SetLocalPosition(mGameScene->GetCamera().GetScreenPos(screenX, 0.9f, 2.f));
	mCoordinateEntity->GetTransform().UpdateMatrices();

	UboCoordinate uboCamera;
	uboCamera.proj = mGameScene->GetCamera().GetProj();
	uboCamera.view = mGameScene->GetCamera().GetView();
	uboCamera.model = mCoordinateEntity->GetTransform().GetWorldMatrix();

	ENQUEUE_RENDER_COMMAND(BuildCoordinate)([this, uboCamera, coordinateModel](VkCommandBuffer commandBuffer)
		{
			mVertexBuffer = VertexBuffer::Builder()
				.SetData(coordinateModel->mVertexs.data())
				.SetSize(coordinateModel->mVertexs.size() * sizeof(VertexCoordinate))
				.Build();

			mIndexBuffer = IndexBuffer::Builder()
				.SetData(coordinateModel->mIndices.data())
				.SetSize(coordinateModel->mIndices.size() * sizeof(uint32))
				.Build();

			UboCoordinate uboCoordinate = uboCamera;
			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboCoordinate)
				.SetSize(sizeof(UboCoordinate))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.SetMultiple(true)
				.Build();


			mPaintDescriptor = RHIDescriptor::Builder()
				.PushShader<CoordinatePaintVertexShader>()
				.PushShader<CoordinatePaintFragmentShader>()
				.SetDescriptorCount(0, 1, true)
				.Build();


			mPaintDescriptor->SetValue("ubo", mUboBuffer);
			mPaintDescriptorMask = mPaintDescriptor->AllocateDescriptor();

			mPaintColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetClearColorValue({ 0.f, 0.f, 0.f, 0.f })
				.SetMultiple(true)
				.Build();

			mPaintPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mPaintDescriptor)
				.AddColorRenderTarget(mPaintColorRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				//.SetPolygonMode(VK_POLYGON_MODE_LINE)
				//.SetCullMode(VK_CULL_MODE_NONE)
				.SetLineWidth(1.f)
				.Build();

			mPaintRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mPaintColorRenderTarget)
				.AddPipeline(mPaintPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mPaintFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mPaintRenderPass)
				.Build();

			RHIDescriptor::Builder quadDescriptorBuilder = RHIDescriptor::Builder();
			mCoordinateQuadVertexShader = quadDescriptorBuilder.AddShader<CoordinateQuadVertexShader>();
			mCoordinateQuadFragmentShader = quadDescriptorBuilder.AddShader<CoordinateQuadFragmentShader>();

			mQuadDescriptor = quadDescriptorBuilder
				.SetDescriptorCount(0, 1, true)
				.Build();

			mSampler = RHISampler::Builder()
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE)
				.Build();

			RefCountPtr<ColorRenderTarget> lastColorRenderTarget = GetTexture(mColorProcess, mColorName);

			mCoordinateQuadFragmentShader->coordinateColor->SetVaule(mPaintColorRenderTarget, mSampler);
			mCoordinateQuadFragmentShader->samplerColor->SetVaule(lastColorRenderTarget, mSampler);

			mQuadDescriptorMask = mQuadDescriptor->AllocateDescriptor();

			mQuadColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build();

			mQuadPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mQuadDescriptor)
				.AddColorRenderTarget(mQuadColorRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mQuadRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mQuadColorRenderTarget)
				.AddPipeline(mQuadPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mQuadFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mQuadRenderPass)
				.Build();
		});
}

void CoordinateProcess::OnTick(float deltaTime)
{
	float screenX = 0.1f * GRenderState->GetHeight() / GRenderState->GetWidth();
	mCoordinateEntity->GetTransform().SetLocalPosition(mGameScene->GetCamera().GetScreenPos(screenX, 0.9f, 2.f));
	mCoordinateEntity->GetTransform().UpdateMatrices();

	UboCoordinate uboCamera;
	uboCamera.proj = mGameScene->GetCamera().GetProj();
	uboCamera.view = mGameScene->GetCamera().GetView();
	uboCamera.model = mCoordinateEntity->GetTransform().GetWorldMatrix();

	ENQUEUE_RENDER_COMMAND(RecordCoordinate)([this, uboCamera](VkCommandBuffer commandBuffer)
		{
			UboCoordinate uboCoordinate = uboCamera;

			mUboBuffer->WriteMemory(&uboCoordinate);


			VkRenderPassBeginInfo renderPassBeginInfo = mPaintFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPaintPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPaintDescriptor->GetPipelineLayout(), 0, 1, &mPaintDescriptorMask.GetDescriptorSet(), 0, nullptr);

			VkBuffer vertexBuffer = mVertexBuffer->GetBuffer();
			VkDeviceSize offsets = 0;
			vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

			VkBuffer indexBuffer = mIndexBuffer->GetBuffer();
			vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

			vkCmdDrawIndexed(commandBuffer, 6, 1, 0, 0, 0);

			vkCmdEndRenderPass(commandBuffer);


			renderPassBeginInfo = mQuadFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mQuadPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mQuadDescriptor->GetPipelineLayout(), 0, 1, &mQuadDescriptorMask.GetDescriptorSet(), 0, nullptr);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);

			GRenderState->SubmitTexture(mQuadColorRenderTarget);
		});
}

void CoordinateProcess::OnResize()
{
	mPaintColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetClearColorValue({ 0.f, 0.f, 0.f, 0.f })
		.SetMultiple(true)
		.Build();

	mPaintFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mPaintColorRenderTarget)
		.SetRenderPass(mPaintRenderPass)
		.Build();

	mQuadColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build();

	mQuadFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mQuadColorRenderTarget)
		.SetRenderPass(mQuadRenderPass)
		.Build();

	RefCountPtr<ColorRenderTarget> lastColorRenderTarget = GetTexture(mColorProcess, mColorName);

	mCoordinateQuadFragmentShader->coordinateColor->SetVaule(mPaintColorRenderTarget, mSampler);
	mCoordinateQuadFragmentShader->samplerColor->SetVaule(lastColorRenderTarget, mSampler);
	mQuadDescriptor->UpdateDescriptor(mQuadDescriptorMask, true);
}

void CoordinateProcess::OnFrameEnd()
{
	if (mColorDirty)
	{
		vkDeviceWaitIdle(GVulkanRHI->GetLogicalDevice());

		RefCountPtr<ColorRenderTarget> lastColorRenderTarget = GetTexture(mColorProcess, mColorName);
		mCoordinateQuadFragmentShader->samplerColor->SetVaule(lastColorRenderTarget, mSampler);
		mQuadDescriptor->UpdateDescriptor(mQuadDescriptorMask, true);

		mColorDirty = false;
	}
}

void CoordinateProcess::UpdateColorRenderTarget(const std::string& colorProcess, const std::string& colorName)
{
	if (mColorProcess == colorProcess && mColorName == colorName)
		return;
	mColorProcess = colorProcess;
	mColorName = colorName;
	mColorDirty = true;
}